# -*- coding: utf-8 -*-
"""
Created on Tue Jun 16 06:29:09 2020

@author: HP
"""


from os.path import isfile
import numpy as np
import pandas as pd
from increasing_closures import lottery_equivalence, time_neutrality, \
    impatience, patience, correlation_neutrality, correlation_aversion, \
    correlation_seeking, stochastic_impatience, stochastic_patience, \
    impatience2
from revealed_python.garp_NOQ import garp_NOQ_find_e

#=============================================================================
#INPUT FILE
#=============================================================================
df = pd.read_csv( 'ITCR_cdata.csv' )
#Get rid of no risk questions
df = df[ (df['pz'] != 66666.) & (df['pw'] != 66666.) ]
ids = sorted( set( df['participantid'] ) )

#=============================================================================
#MODES AND COLUMNS
#=============================================================================
imple_list = [impatience, lottery_equivalence]
tests_dict = {'def': [], \
              'le': [lottery_equivalence], \
              'tn': [time_neutrality], \
              'imp': [impatience], \
              'pat': [patience], \
              'cn': [correlation_neutrality], \
              'ca': [correlation_aversion], \
              'imple': [impatience, lottery_equivalence], \
              'imple_ca': imple_list + [correlation_aversion], \
              'imple_cs': imple_list + [correlation_seeking], \
              'imple_si': imple_list + [stochastic_impatience], \
              'imple_sp': imple_list + [stochastic_patience], \
              'imple_ca_si': imple_list + [correlation_aversion,stochastic_impatience], \
              'imple_ca_sp': imple_list + [correlation_aversion,stochastic_patience], \
              'imple_cs_si': imple_list + [correlation_seeking,stochastic_impatience], \
              'imple_cs_sp': imple_list + [correlation_seeking,stochastic_patience]}

columns = list(tests_dict.keys())

#-------------------
#manage results file
#-------------------
results_file = 'rp_tests_by_participants.xlsx'
results_df = None
if isfile( results_file ):
    results_df = pd.read_excel( results_file, sheet_name = 0, index_col = 0 )
else:
    #Make a new results df
    results_df = pd.DataFrame( np.nan, index = ids, columns = columns )
    results_df.index.name = 'participantid'
#-------------------------
#for each id run the test.
#-------------------------

change_flag = False
for pid in ids:
    print('current id: {0}'.format(pid))
    cur_df = df[ df['participantid'] == pid ]
    C = cur_df[['x','y','z','w']].to_numpy()
    P = cur_df[['px','py','pz','pw']].to_numpy()
    for test_name in tests_dict:
        #Run the test if we don't already have the results.
        if np.isnan( results_df.loc[pid,test_name] ):
            e = garp_NOQ_find_e(C,P,Ilist= tests_dict[test_name])
            results_df.loc[pid,test_name] = e
            change_flag = True
    if change_flag:
        change_flag = False      
        results_df.to_excel(results_file)
        
#=============================================================================
#DATA SPLITTING
#=============================================================================

df1 = df[ ['participantid','px','py','x','y'] ].copy()
df2 = df[ ['participantid','pz','pw','z','w'] ].copy()
df1.columns = [ 'participantid', 'p1', 'p2', 'q1', 'q2' ]
df2.columns = [ 'participantid', 'p1', 'p2', 'q1', 'q2' ]

df = pd.concat([df1,df2], axis=0, ignore_index = True)
df = df.sort_values(by='participantid')

if 'imp2split' not in results_df.columns:
    results_df['imp2split'] = np.nan

change_flag = False
for pid in ids:
    print('current id: {0}'.format(pid))
    cur_df = df[ df['participantid'] == pid ]
    C = cur_df[['q1','q2']].to_numpy()
    P = cur_df[['p1','p2']].to_numpy()
    if np.isnan( results_df.loc[pid,'imp2split'] ):
        e = garp_NOQ_find_e(C,P,Ilist= [impatience2])
        results_df.loc[pid,'imp2split'] = e
        change_flag = True
    if change_flag:
        change_flag = False      
        results_df.to_excel(results_file)
        
results_df[ 'imple_imp2split' ] = results_df[ ['imple','imp2split'] ].min(1)
results_df.to_excel(results_file)
        
#Classify participants
#-------------------
#manage results file
#-------------------
results_file = 'rp_tests_by_participants.xlsx'
results_df = None
if isfile( results_file ):
    results_df = pd.read_excel( results_file, sheet_name = 0, index_col = 0 )
else:
    #Make a new results df
    raise Exception('Result file should exist.')
    
    
#If results_df has corners or edges vars then drop them.

if 'corners' in results_df.columns:
    results_df.drop('corners',axis=1, inplace=True)
if 'edges' in results_df.columns:
    results_df.drop('edges',axis=1, inplace=True)
#-------------------------
#Add a corner flag. Equals 1 if all rows have a 0. Otherwise returns 1.
#-------------------------

df['edges'] = 0
df['corners'] = 0
zeros_count = (df[['x','y','z','w']] == 0.).sum(1)
df.loc[ zeros_count >= 1, 'edges' ] = 1
df.loc[ zeros_count >= 3, 'corners' ] = 1

char_df = df[['participantid','edges','corners']].groupby('participantid').sum()
results_df = results_df.join(char_df)
results_df.to_excel(results_file)


